package com.study.regression

import com.study.classification.Iris
import org.apache.spark.ml.Pipeline
import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.feature.{IndexToString, StringIndexer, VectorIndexer}
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.regression.{DecisionTreeRegressionModel, DecisionTreeRegressor}
import org.apache.spark.sql.SparkSession

/**
 * 回归-决策树（使用鸢尾花数据集）
 *
 * @author stephen
 * @date 2019-08-27 18:19
 */

case class Iris(features: org.apache.spark.ml.linalg.Vector, label: String)

object DecisionTreeRegressorIrisDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")

    import spark.implicits._

    val path = this.getClass.getClassLoader.getResource("data/iris.data").getPath
    val data = spark.read
      .textFile(path)
      .map(_.split(","))
      .map(p =>
        Iris(Vectors.dense(p(0).toDouble, p(1).toDouble, p(2).toDouble, p(3).toDouble), p(4)))
      .toDF()
    data.createOrReplaceTempView("iris")
    val df = spark.sql("select * from iris")

    // 分别获取标签列和特征列，进行索引，并进行了重命名。
    val labelIndexer = new StringIndexer()
      .setInputCol("label")
      .setOutputCol("indexedLabel")
      .fit(df)
    val featureIndexer = new VectorIndexer()
      .setInputCol("features")
      .setOutputCol("indexedFeatures")
      .setMaxCategories(4)
      .fit(df)
    // 设置一个labelConverter，目的是把预测的类别重新转化成字符型的。
    val labelConverter = new IndexToString()
      .setInputCol("prediction")
      .setOutputCol("predictedLabel")
      .setLabels(labelIndexer.labels)
    // 把数据集随机分成训练集和测试集，其中训练集占70%。
    val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
    // 训练决策树模型,这里我们可以通过setter的方法来设置决策树的参数，也可以用ParamMap来设置。
    // 具体可以设置的参数可以通过explainParams()来获取。
    val dtRegressor = new DecisionTreeRegressor()
      .setLabelCol("indexedLabel")
      .setFeaturesCol("indexedFeatures")
    // 在pipeline中进行设置
    val pipelineRegressor = new Pipeline()
      .setStages(Array(labelIndexer, featureIndexer, dtRegressor, labelConverter))
    // 训练决策树模型
    val modelRegressor = pipelineRegressor.fit(trainingData)
    //进行预测
    val predictionsRegressor = modelRegressor.transform(testData)
    // 查看部分预测结果
    predictionsRegressor.select("predictedLabel", "label", "features").show(20)

    // 评估决策树回归模型
    val evaluatorRegressor = new RegressionEvaluator()
      .setLabelCol("indexedLabel")
      .setPredictionCol("prediction")
      .setMetricName("rmse")
    // 得到均方根误差
    val rmse = evaluatorRegressor.evaluate(predictionsRegressor)
    println("Root Mean Squared Error (RMSE) on test data = " + rmse)

    val treeModelRegressor = modelRegressor
      .stages(2)
      .asInstanceOf[DecisionTreeRegressionModel]
    println("Learned regression tree model:\n" + treeModelRegressor.toDebugString)
  }
}
